{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Finding the best open-source embedding model for your RAG application"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this tutorial, we explore the process of selecting the most suitable open-source embedding model for your Retrieval-Augmented Generation (RAG) application. \n",
    "\n",
    "We evaluate three open-source embedding models: [nomic-embed-text](https://ollama.com/library/nomic-embed-text), [bge-m3](https://ollama.com/library/bge-m3/blobs/daec91ffb5dd), and [mxbai-embed-large](https://ollama.com/library/mxbai-embed-large/blobs/819c2adf5ce6).\n",
    " \n",
    "To facilitate the execution of these models and the generation of embeddings, we will use the following tools:\n",
    "\n",
    "- [Ollama](https://ollama.com/), a platform that provides access to a variety of open-source Large Language Models (LLMs).\n",
    "- [pgai](https://github.com/timescale/pgai), an open-source extension that seamlessly integrates LLM workflows, such as embedding creation and management, directly into your PostgreSQL database.\n",
    "\n",
    "The evaluation process involves:\n",
    "\n",
    "1. Setting up a test environment with Ollama and PostgreSQL\n",
    "2. Loading Paul Graham's essays as our test dataset\n",
    "3. Generating embeddings using different models\n",
    "4. Creating diverse test questions across multiple categories\n",
    "5. Evaluating each model's retrieval performance"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Environment Setup\n",
    "\n",
    "Before you begin, execute the `install.sh` shell script to set up your environment. This script will install all necessary Docker containers required for the environment.\n",
    "\n",
    "**Note for macOS Users:**\n",
    "  \n",
    "To run the script, first make it executable by running the command: `chmod +x install.sh`. Then, execute the script with: `./install.sh`.\n",
    "\n",
    "Let's install the necessary Python libraries for this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install pandas psycopg2-binary Jinja2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's define the environment variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATABASE_CONNECTION_STRING=\"postgres://postgres:postgres@localhost:5432/postgres\"\n",
    "OLLAMA_HOST=\"http://ollama:11434\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset Ingestion\n",
    "\n",
    "Let's setup the PostgreSQL database and install [pgai](https://github.com/timescale/pgai)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import psycopg2\n",
    "\n",
    "def connect_db():\n",
    "    return psycopg2.connect(DATABASE_CONNECTION_STRING)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "with connect_db() as conn:\n",
    "    with conn.cursor() as cur:\n",
    "        cur.execute(\"CREATE EXTENSION IF NOT EXISTS ai CASCADE;\")\n",
    "\n",
    "    with conn.cursor() as cur:\n",
    "        cur.execute(\"\"\"\n",
    "            CREATE TABLE IF NOT EXISTS essays (\n",
    "                id BIGINT PRIMARY KEY GENERATED BY DEFAULT AS IDENTITY,\n",
    "                title TEXT NOT NULL,\n",
    "                date TEXT,\n",
    "                text TEXT NOT NULL\n",
    "            );\n",
    "          \"\"\") "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Loading the Dataset\n",
    "\n",
    "We'll now load [Paul Graham's essays](https://huggingface.co/datasets/sgoel9/paul_graham_essays) into our PostgreSQL database. The dataset consists of 215 essays with their titles, dates, and full text content. We'll verify the successful ingestion by displaying the first essay from the database."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with connect_db() as conn:\n",
    "    with conn.cursor() as cur:\n",
    "        # Load Paul Graham's essays dataset into the 'essays' table\n",
    "        cur.execute(\"\"\"\n",
    "            SELECT ai.load_dataset(\n",
    "                    'sgoel9/paul_graham_essays', \n",
    "                    table_name => 'essays', \n",
    "                    if_table_exists => 'append');\n",
    "        \"\"\")\n",
    "    \n",
    "    with conn.cursor() as cur:\n",
    "        # Fetch and print the first row from the 'essays' table to verify the data\n",
    "        cur.execute(\"SELECT * FROM essays LIMIT 1;\")\n",
    "        print(cur.fetchone())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generating Embeddings\n",
    "\n",
    "pgai makes generating embeddings for various models incredibly straightforward with its [pgai Vectorizer](https://github.com/timescale/pgai/blob/main/docs/vectorizer.md). With just a single SQL command, you can effortlessly configure a vectorizer to automatically create and update embeddings from any chosen embedding model.\n",
    "\n",
    "We create vectorizers for each embedding model. Each vectorizer will:\n",
    "\n",
    "- Process text using the same chunking strategy (512 characters with 50 character overlap)\n",
    "- Use consistent text formatting across all models\n",
    "- Generate embeddings with model-specific dimensions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_vectorizer(embedding_model, embeddings_dimensions):\n",
    "    embeddings_view_name = f\"{'essays'}{'_'}{embedding_model.replace('-','_')}{'_'}{'embeddings'}\"\n",
    "\n",
    "    with connect_db() as conn:\n",
    "        with conn.cursor() as cur:\n",
    "            cur.execute(\"\"\"\n",
    "                SELECT ai.create_vectorizer(\n",
    "                'essays'::regclass,\n",
    "                destination => %s,\n",
    "                loading => ai.loading_column('text'),\n",
    "                embedding => ai.embedding_ollama(%s, %s),\n",
    "                chunking => ai.chunking_recursive_character_text_splitter(512, 50),\n",
    "                formatting => ai.formatting_python_template('title: $title $chunk')\n",
    "            );\n",
    "            \"\"\", (embeddings_view_name, embedding_model, embeddings_dimensions, )\n",
    "            )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "EMBEDDING_MODELS = [\n",
    "    {'name':'mxbai-embed-large', 'dimensions': 1024},\n",
    "    {'name':'nomic-embed-text','dimensions': 768},\n",
    "    {'name':'bge-m3','dimensions': 1024},\n",
    "] \n",
    "\n",
    "for model in EMBEDDING_MODELS:\n",
    "    create_vectorizer(model['name'], model['dimensions'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The vectorizers will take some time to complete the embedding generation. Use the following command to monitor their progress:  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with connect_db() as conn:\n",
    "    with conn.cursor() as cur:\n",
    "        cur.execute(\"SELECT * FROM ai.vectorizer_status;\")\n",
    "\n",
    "        for row in cur.fetchall():\n",
    "            print(f\"Vectorizer ID: {row[0]}, Embedding Table: {row[2]}, Pending Items: {row[4]}\")\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Embeddings Evaluation\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluation Parameters\n",
    "\n",
    "To ensure a fair comparison, we'll establish consistent evaluation parameters:\n",
    "\n",
    "- Number of text chunks to evaluate\n",
    "- Questions per chunk across different categories\n",
    "- Number of top results to consider (K)\n",
    "- Distribution of question types to test different aspects of embedding quality"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "NUM_CHUNKS = 20\n",
    "NUM_QUESTIONS_PER_CHUNK = 20\n",
    "TOP_K = 10\n",
    "\n",
    "QUESTION_DISTRIBUTION = {\n",
    "    'short': 4,\n",
    "    'long': 4,\n",
    "    'direct': 4,\n",
    "    'implied': 4,\n",
    "    'unclear': 4\n",
    "}\n",
    "\n",
    "assert sum(QUESTION_DISTRIBUTION.values()) == NUM_QUESTIONS_PER_CHUNK"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluation Chunks\n",
    "\n",
    "We select 20 random chunks from one of the embeddings view since all three models use the same chunks. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd \n",
    "\n",
    "evaluation_chunks = []\n",
    "\n",
    "with connect_db() as conn:\n",
    "    with conn.cursor() as cur:\n",
    "        cur.execute(\"\"\"\n",
    "                SELECT id, chunk_seq, chunk, title \n",
    "                FROM essays_nomic_embed_text_embeddings \n",
    "                ORDER BY RANDOM() \n",
    "                LIMIT %s\n",
    "            \"\"\", (NUM_CHUNKS,))\n",
    "        \n",
    "        for row in cur.fetchall():\n",
    "            evaluation_chunks.append({\n",
    "                'id': row[0],\n",
    "                'chunk_seq': row[1],\n",
    "                'chunk': row[2],\n",
    "                'title': row[3]\n",
    "            })\n",
    "\n",
    "pd.DataFrame(evaluation_chunks).to_csv('./chunks.csv')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluation Question Generation\n",
    "\n",
    "We'll generate diverse questions for each text chunk across five categories:\n",
    "\n",
    "1. **Short questions**: Simple, direct queries under 10 words\n",
    "2. **Long questions**: Detailed questions requiring comprehensive understanding\n",
    "3. **Direct questions**: Questions about explicit information\n",
    "4. **Implied questions**: Questions requiring contextual understanding\n",
    "5. **Unclear questions**: Ambiguous queries to test robustness\n",
    "\n",
    "Each category tests different aspects of the embedding models' capabilities."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_questions_by_question_type(chunk, question_type, num_questions):\n",
    "    prompts = {\n",
    "        'short': \"Generate {count} short, simple questions about this text. Questions should be direct, under 10 words\",\n",
    "        'long': \"Generate {count} detailed, comprehensive questions about this text. Include specific details:\",\n",
    "        'direct': \"Generate {count} questions that directly ask about explicit information in this text\",\n",
    "        'implied': \"Generate {count} questions that require understanding context and implications of the text:\",\n",
    "        'unclear': \"Generate {count} vague, ambiguous questions about the general topic of this text:\"\n",
    "    }\n",
    "\n",
    "    prompt = prompts[question_type].format(count=num_questions) + f\"\\n\\nText: {chunk}\"\n",
    "\n",
    "    system_instructions = \"\"\"\n",
    "        Generate different types of questions about the given text following the prompt provided. \n",
    "        Each question must be on a new line. Do not include empty lines or blank questions.\n",
    "    \"\"\"\n",
    "\n",
    "    with connect_db() as conn:\n",
    "        with conn.cursor() as cur:\n",
    "            cur.execute(\"\"\"\n",
    "                SELECT ai.ollama_generate(\n",
    "                    'llama3.2',\n",
    "                    %s,\n",
    "                    system_prompt=>%s, \n",
    "                    host=>%s\n",
    "                )->>'response';\n",
    "            \"\"\",(prompt, system_instructions, OLLAMA_HOST))\n",
    "\n",
    "            generated_questions = [q.strip() for q in cur.fetchone()[0].split(\"\\n\") if q.strip()]\n",
    "            print(f\"Number of questions generated for {question_type}: {len(generated_questions)}\")\n",
    "            return generated_questions "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluation_questions = []\n",
    "\n",
    "for i, chunk in enumerate(evaluation_chunks, 1):\n",
    "    print(f\"Processing chunk {i}/{len(evaluation_chunks)}\")\n",
    "\n",
    "    for question_type, count in QUESTION_DISTRIBUTION.items():\n",
    "        questions = generate_questions_by_question_type(chunk['chunk'], question_type, count)\n",
    "\n",
    "        for q in questions:\n",
    "            evaluation_questions.append({\n",
    "                'question': q,\n",
    "                'source_chunk_id': chunk['id'],\n",
    "                'source_chunk_seq': chunk['chunk_seq'],\n",
    "                'question_type': question_type,\n",
    "                'chunk': chunk['chunk']\n",
    "            })\n",
    "\n",
    "print(\"Generated questions in total:\", len(evaluation_questions))\n",
    "\n",
    "pd.DataFrame(evaluation_questions).to_csv('./generated_questions.csv')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Embedding Model Evaluation\n",
    "\n",
    "The evaluation process involves:\n",
    "\n",
    "1. For each model:\n",
    "\n",
    "   - Converting questions to embeddings\n",
    "   - Performing similarity search against chunk embeddings\n",
    "   - Checking if the source chunk appears in `top-K` results\n",
    "\n",
    "2. Calculating performance metrics:\n",
    "\n",
    "   - Overall accuracy\n",
    "   - Per-question-type accuracy\n",
    "   - Detailed success/failure analysis\n",
    "\n",
    "This comprehensive evaluation will help identify each model's strengths and weaknesses."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "OLLAMA_HOST = os.environ[\"OLLAMA_HOST\"]\n",
    "\n",
    "def vector_similarity_search(embeddings_view, embedding_model, question):\n",
    "    with connect_db() as conn:\n",
    "        with conn.cursor() as cur:\n",
    "            cur.execute(f\"\"\"\n",
    "                SELECT id, chunk_seq \n",
    "                FROM {embeddings_view} \n",
    "                ORDER BY embedding <=> ai.ollama_embed(%s, %s, host => %s)\n",
    "                LIMIT %s;\n",
    "            \"\"\", (embedding_model, question, OLLAMA_HOST, TOP_K,)\n",
    "            )\n",
    "\n",
    "            return cur.fetchall()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_embedding_models():\n",
    "    evaluation_results = []\n",
    "    detailed_results = []\n",
    "\n",
    "    for model in EMBEDDING_MODELS:\n",
    "        print(f\"Evaluating {model['name']}...\")\n",
    "\n",
    "        embeddings_view = f\"{'essays'}{'_'}{model['name'].replace('-','_')}{'_'}{'embeddings'}\"\n",
    "        scores = []\n",
    "\n",
    "        for q in evaluation_questions:\n",
    "            vector_search_results = vector_similarity_search(embeddings_view, model['name'], q['question'])\n",
    "            found = any(\n",
    "                row[0] == q['source_chunk_id'] and row[1]== q['source_chunk_seq'] \n",
    "                for row in vector_search_results\n",
    "            )\n",
    "\n",
    "            scores.append(1 if found else 0)\n",
    "\n",
    "            detailed_results.append({\n",
    "                'model': model['name'],\n",
    "                'question': q['question'],\n",
    "                'question_type': q['question_type'],\n",
    "                'source_chunk_id': q['source_chunk_id'],\n",
    "                'source_chunk_seq': q['source_chunk_seq'],\n",
    "                'found_correct_chunk': found,\n",
    "                'num_results': len(vector_search_results)\n",
    "            })\n",
    "\n",
    "        evaluation_results.append({\n",
    "            'model': model['name'],\n",
    "            'overall_accuracy': sum(scores) / len(scores),\n",
    "            'by_type': {\n",
    "                q_type: sum(scores[i] for i, q in enumerate(evaluation_questions) \n",
    "                            if q['question_type'] == q_type) / QUESTION_DISTRIBUTION[q_type] / NUM_CHUNKS\n",
    "                for q_type in QUESTION_DISTRIBUTION.keys()\n",
    "            }\n",
    "        })\n",
    "    \n",
    "    pd.DataFrame(detailed_results).to_csv('./evaluation_data/detailed_results.csv')\n",
    "    return evaluation_results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_results_table(evaluation_results):\n",
    "    # Create lists to store the data\n",
    "    rows = []\n",
    "    \n",
    "    # Process each model's results\n",
    "    for result in evaluation_results:\n",
    "        row = {\n",
    "            'Model': result['model'],\n",
    "            'Overall Accuracy': f\"{result['overall_accuracy']:.2%}\",\n",
    "        }\n",
    "        # Add accuracies for each question type\n",
    "        for q_type, acc in result['by_type'].items():\n",
    "            row[q_type.capitalize()] = f\"{acc:.2%}\"\n",
    "        \n",
    "        rows.append(row)\n",
    "    \n",
    "    # Create DataFrame\n",
    "    df = pd.DataFrame(rows)\n",
    "    \n",
    "    # Reorder columns to put Overall Accuracy after Model\n",
    "    columns = ['Model', 'Overall Accuracy'] + [col for col in df.columns if col not in ['Model', 'Overall Accuracy']]\n",
    "    df = df[columns]\n",
    "    \n",
    "    # Display the table\n",
    "    return df.style.set_properties(**{\n",
    "        'text-align': 'center',\n",
    "        'border': '1px solid black',\n",
    "        'padding': '8px'\n",
    "    }).set_table_styles([\n",
    "        {'selector': 'th', 'props': [\n",
    "            ('background-color', 'black'),\n",
    "            ('text-align', 'center'),\n",
    "            ('padding', '8px'),\n",
    "            ('border', '1px solid black')\n",
    "        ]},\n",
    "        {'selector': 'caption', 'props': [\n",
    "            ('text-align', 'center'),\n",
    "            ('font-weight', 'bold'),\n",
    "            ('font-size', '1.1em'),\n",
    "            ('padding', '8px')\n",
    "        ]}\n",
    "    ]).set_caption('Embedding Models Evaluation Results')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Display the results\n",
    "evaluation_results = evaluate_embedding_models()\n",
    "results_table = create_results_table(evaluation_results)\n",
    "display(results_table)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pythonenv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
